# 1. 环境与LLM初始化
import os
import json
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI

load_dotenv()
llm = ChatOpenAI(
    model_name="qwen-plus",
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    temperature=0
)

# 2. 定义结构化工具
from langchain_core.tools import tool


@tool(
    description="查询气温，参数格式：城市,日期（YYYY-MM-DD），示例：北京,2024-10-15"
)
def WeatherTool(query: str) -> str:
    try:
        city, date = query.split(",")
        city, date = city.strip(), date.strip()
        weather_data = {("北京", "2024-10-15"): "23℃", ("上海", "2024-10-15"): "26℃"}
        return f"{city} {date}：{weather_data.get((city, date), '暂无数据')}"
    except ValueError:
        return "天气参数格式错误！示例：北京,2024-10-15"


@tool(
    description="预订酒店，参数为JSON字符串（需包含city、date、room_type、guests），示例：{\"city\":\"北京\",\"date\":\"2024-10-15\",\"room_type\":\"大床房\",\"guests\":2}"
)
def HotelTool(json_str: str) -> str:
    try:
        params = json.loads(json_str)
        required = ["city", "date", "room_type", "guests"]
        if not all(k in params for k in required):
            return f"缺少必填参数！需包含：{', '.join(required)}"
        return f"已预订 {params['city']} {params['date']} {params['room_type']}（{params['guests']}人），订单号：H{hash(json_str) % 1000000}"
    except json.JSONDecodeError:
        return "酒店参数JSON格式错误！示例：{\"city\":\"北京\",\"date\":\"2024-10-15\",\"room_type\":\"大床房\",\"guests\":2}"


tools = [WeatherTool, HotelTool]

# 3. 初始化记忆（使用简单的字典存储）
chat_history = []


# 4. 创建带记忆的Agent
def agent_with_memory(user_input: str):
    """带记忆的Agent调用"""
    global chat_history

    # 构建包含历史对话的系统提示
    history_text = ""
    if chat_history:
        history_text = "\n".join([f"用户：{h['user']}\n助手：{h['assistant']}" for h in chat_history[-3:]])  # 只保留最近3轮对话

    system_prompt = f"""你是一个智能助手，可以使用以下工具：

工具列表：
- WeatherTool: 查询气温，参数格式：城市,日期（YYYY-MM-DD），示例：北京,2024-10-15
- HotelTool: 预订酒店，参数为JSON字符串（需包含city、date、room_type、guests），示例：{{"city":"北京","date":"2024-10-15","room_type":"大床房","guests":2}}

历史对话：
{history_text}

重要规则：
1. 如果用户查询天气，使用WeatherTool，参数格式：城市,日期
2. 如果用户预订酒店，使用HotelTool，参数为JSON格式
3. 如果用户没有提供完整的酒店预订信息（如城市、日期），请从历史对话中提取最近的城市和日期
4. 如果历史对话中没有相关信息，请询问用户补充

请根据用户输入，决定是否需要调用工具。如果需要调用工具，请按以下格式：
Action: 工具名
Action Input: 参数

如果不需要工具，直接回答用户问题。

用户输入：{user_input}"""

    # 调用LLM
    response = llm.invoke(system_prompt)

    # 检查是否需要调用工具
    if "Action:" in response.content:
        lines = response.content.split('\n')
        action_line = None
        action_input_line = None

        for line in lines:
            if line.startswith("Action:"):
                action_line = line
            elif line.startswith("Action Input:"):
                action_input_line = line

        if action_line and action_input_line:
            action = action_line.split("Action:")[1].strip()
            action_input = action_input_line.split("Action Input:")[1].strip()

            # 调用相应的工具
            if action == "WeatherTool":
                result = WeatherTool.invoke(action_input)
            elif action == "HotelTool":
                result = HotelTool.invoke(action_input)
            else:
                result = "未知工具"

            # 保存到历史记录
            chat_history.append({
                "user": user_input,
                "assistant": f"工具调用结果：{result}"
            })

            return f"工具调用结果：{result}"

    # 保存到历史记录
    chat_history.append({
        "user": user_input,
        "assistant": response.content
    })

    return response.content


# 5. 测试多轮对话
if __name__ == "__main__":
    print("===== 测试1：查天气 =====")
    res1 = agent_with_memory("查询北京2024-10-15的天气")
    print("Agent回答：", res1, "\n")

    print("===== 测试2：订酒店（复用参数） =====")
    res2 = agent_with_memory("预订酒店，大床房，2人")
    print("Agent回答：", res2, "\n")

    print("===== 测试3：查新城市天气 =====")
    res3 = agent_with_memory("查询上海2024-10-15的天气")
    print("Agent回答：", res3, "\n")

    print("===== 测试4：再次订酒店（复用新参数） =====")
    res4 = agent_with_memory("预订酒店，双床房，3人")
    print("Agent回答：", res4)
